import json
from pathlib import Path
import pandas as pd


def convert2submit(test_file: Path, prediction_file: Path, save_path: Path):
    pred_label_list = []

    for line in open(prediction_file, "r"):
        prediction_data = json.loads(line)

        pred_label = prediction_data["predict"]
        pred_label_list.append(pred_label)

    test_data = json.load(open(test_file, "r"))
    save_data = []
    for i, example in enumerate(test_data):
        example["predict"] = pred_label_list[i]
        save_data.append(example)

    df = pd.DataFrame(save_data)

    df.to_csv(save_path, index=None, encoding="utf-8-sig")


if __name__ == "__main__":
    # test_file = "data/demo_test.json"
    # prediction_file = "data/demo_pred.jsonl"
    # save_path = "submit.csv"
    # 测试集
    test_file = "/root/LLaMA-Factory/data/mire/test1/test1.json"
    prediction_file_list = [
        "/root/autodl-fs/generated_predictions.jsonl",
    ]
    # # 训练集
    # test_file = "/home/lyc/TNTprojectz/LLaMA-Factory/data/mire/train0/train0.json"
    # prediction_file_list = [
    #     '/share/saves/qwen2_vl-2b/full/sft-infer-ckpt120/generated_predictions.jsonl',
    #     '/share/saves/qwen2_vl-2b/lora/sft-infer-ckpt300/generated_predictions.jsonl',
    #     '/share/saves/qwen2_vl-7b/lora/sft-infer-ckpt240/generated_predictions.jsonl'
    # ]
    for prediction_file in prediction_file_list:
        # save_path = "submit.csv"
        # save_path = prediction_file.replace("generated_predictions.jsonl", "submit.csv").replace('/share/saves/', '').replace('/', '-')
        save_path = prediction_file.replace("generated_predictions.jsonl", "submit.csv").replace('/root/autodl-fs/', '').replace('/', '-')
        save_path = f'./submit/{save_path}'
        print(save_path)
        convert2submit(test_file, prediction_file, save_path)

# end main
